import os
import torch
import numpy as np
import random
from ultralytics import YOLO
from utils import prepare_dataset

# 设置随机种子，保证实验可重复性
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# 定义模型参数
class Config:
    # 数据集相关参数
    annotations = 'data/annotations'
    images = 'data/images'
    output = 'data/prepared'
    class_names = ['with_mask', 'without_mask', 'mask_weared_incorrect']
    
    # 模型相关参数
    model = 'yolov8n.pt'       # 使用的模型：yolov8n.pt, yolov8s.pt 等
    weights = ''               # 初始权重路径，留空则使用预训练模型
    epochs = 100               # 训练轮数
    batch_size = 16            # 批次大小
    img_size = 640             # 图像大小
    hyp = 'data/hyp.yaml'      # 超参数配置路径
    
    # 其他配置
    skip_data_prep = True     # 是否跳过数据准备
    seed = 42                  # 随机种子
    project = 'runs'           # 保存项目路径
    name = 'train'             # 训练名称
    exist_ok = True            # 是否允许覆盖现有实验目录

# 主函数
def main():
    # 获取配置
    cfg = Config()
    
    # 设置随机种子
    set_seed(cfg.seed)
    
    # 处理数据集
    if cfg.skip_data_prep:
        print("跳过数据准备步骤...")
        data_cfg = os.path.join(cfg.output, 'dataset.yaml')
        if not os.path.exists(data_cfg):
            print(f"警告: 未找到数据集配置文件 {data_cfg}，将创建新配置")
            from utils import create_dataset_config
            data_cfg = create_dataset_config(
                cfg.output,
                cfg.class_names,
                os.path.join(cfg.output, 'train', 'images'),
                os.path.join(cfg.output, 'val', 'images'),
                os.path.join(cfg.output, 'test', 'images')
            )
    else:
        # 准备数据集
        data_cfg = prepare_dataset(cfg.annotations, cfg.images, cfg.output, cfg.class_names)
    
    # 训练模型 (使用ultralytics API)
    print("Starting model training...")
    
    # 加载模型
    if cfg.weights and os.path.exists(cfg.weights):
        model = YOLO(cfg.weights)
    else:
        model = YOLO(cfg.model)
    
    # 使用新版API训练模型
    model.train(
        data=data_cfg,
        epochs=cfg.epochs,
        batch=cfg.batch_size,
        imgsz=cfg.img_size,
        project=cfg.project,
        name=cfg.name,
        exist_ok=cfg.exist_ok
    )
    
    print(f"Training complete. Model saved at {cfg.project}/{cfg.name}")

if __name__ == "__main__":
    main() 